package com.ht.modules.component.ai.session.service;

import com.ht.modules.component.ai.session.entity.SessionInfo;
import com.ht.modules.component.ai.session.repository.SessionRepository;
import dev.langchain4j.data.message.*;
import org.apache.logging.log4j.util.Strings;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Example;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class SessionDbService {
    @Autowired
    private SessionRepository sessionRepository;

    /**
     * 根据 messageId 查询消息
     *
     * @param messageId 消息ID
     * @return 消息
     */
    public List<ChatMessage> getMessages(String messageId) {
        List<SessionInfo> messages = this.list(messageId);
        if (messages.isEmpty()) {
            return new ArrayList<>();
        }
        List<ChatMessage> chatMessageList = new ArrayList<>();
        for (SessionInfo message : messages) {
            ChatMessageType type = message.getType();
            String data = message.getMessage();
            ChatMessage chatMessage = switch (type) {
                case SYSTEM -> SystemMessage.systemMessage(data);
                case USER -> UserMessage.userMessage(data);
                case AI -> AiMessage.aiMessage(data);
                case TOOL_EXECUTION_RESULT -> ToolExecutionResultMessage.toolExecutionResultMessage(null, null, data);
                case CUSTOM -> {
                    Map<String, Object> param = new HashMap<>();
                    param.put("text", data);
                    yield CustomMessage.customMessage(param);
                }
            };
            chatMessageList.add(chatMessage);
        }
        return chatMessageList;
    }

    /**
     * 保存消息
     *
     * @param sessionId       会话ID
     * @param chatMessageList 消息
     */
    public void updateMessages(String sessionId, List<ChatMessage> chatMessageList) {
        List<SessionInfo> sessionInfoList = new ArrayList<>();
        for (ChatMessage chatMessage : chatMessageList) {
            SessionInfo message = SessionInfo.builder()
                    .sessionId(sessionId)
                    .type(chatMessage.type())
                    .build();
            String data = "";
            switch (chatMessage.type()) {
                case SYSTEM -> {
                    SystemMessage systemMessage = (SystemMessage) chatMessage;
                    data = systemMessage.text();
                }
                case USER -> {
                    UserMessage userMessage = (UserMessage) chatMessage;
                    data = userMessage.singleText();
                }
                case AI -> {
                    AiMessage aiMessage = (AiMessage) chatMessage;
                    data = aiMessage.text();
                }
                case TOOL_EXECUTION_RESULT -> {
                    ToolExecutionResultMessage toolExecutionResultMessage = (ToolExecutionResultMessage) chatMessage;
                    data = toolExecutionResultMessage.text();
                }
                case CUSTOM -> {
                    CustomMessage customMessage = (CustomMessage) chatMessage;
                    data = customMessage.attributes().get("text").toString();
                }
            }
            message.setMessage(data);
            message.setDataTime(String.valueOf(System.currentTimeMillis()));
            if (Strings.isBlank(message.getMessage())) {
                continue;
            }
            sessionInfoList.add(message);
        }
        this.deleteMessages(sessionId);
        if (sessionInfoList.isEmpty()) {
            return;
        }
        this.saveBatch(sessionInfoList);
    }


    /**
     * 根据 messageId 删除消息
     *
     * @param messageId 消息ID
     */
    public void deleteMessages(String messageId) {
        List<SessionInfo> messages = this.list(messageId);
        if (messages.isEmpty()) {
            return;
        }
        sessionRepository.deleteAll(messages);
    }


    /**
     * 根据 messageId 查询消息
     *
     * @return 消息
     */
    private List<SessionInfo> list(String messageId) {
        return sessionRepository.findAll(Example.of(SessionInfo.builder().sessionId(messageId).build()));
    }

    /**
     * 批量保存消息
     *
     * @param sessionInfoList 消息
     */
    private void saveBatch(List<SessionInfo> sessionInfoList) {
        sessionRepository.saveAll(sessionInfoList);
    }
}
